import argparse
import os
from sklearn.cluster import KMeans
from torch_geometric.transforms import NormalizeFeatures
from Metric import Metric
from Output import
from ProcessData import DependenceGraph, NoLabelDependenceGraph, FileVectorDependenceGraph
from ProcessData.CodeVectorDependenceGraph import codevectorDependenceGraph
from Utils import ReadConfig
from Output import Vote


def main():
    configpath = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "config.ini")
    parser = argparse.ArgumentParser(description="ReadConfig")
    parser.add_argument(
        "-c",
        "--config",
        type=str,
        default=configpath,
        help="this is the path of Config file"
    )

    root, P, model, outfilepath = ReadConfig(configpath)
    a = 2
    if a == 1:
        dataset = FileVectorDependenceGraph(root, transform=NormalizeFeatures())
    if a == 2:
        dataset = codevectorDependenceGraph(root, transform=NormalizeFeatures())
    if a == 3:
        dataset = NoLabelDependenceGraph(root, transform=NormalizeFeatures())
    if a == 4:
        dataset = DependenceGraph(root, transform=NormalizeFeatures())
    data = dataset[0]

    z = data.x
    kmeans_input = z.cpu().numpy()
    kmeans = KMeans(n_clusters=14, random_state=0).fit(kmeans_input)
    preds = kmeans.predict(kmeans_input)
    if a==1:
        result2rsf_file(root, preds, outfilepath)

    if a==2:
        new_pred=Vote(dataset.dic, "D:\module-reverse-by-gnn\mdgdata\\raw\\filedefsymbol.json", preds, outfilepath)
    print(Metric(P, outfilepath))


if __name__ == "__main__":
    main()
